# train.py

import os
import torch
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
import json

from datasets.dataset import TicketDataset
from models.MPDnnModel import PricePredictor
from utils.preprocess import preprocess
from utils.concat_csv_files import concat_csv_files

# ========== 配置参数 ==========
TRAIN_DIR = "data/train"
TEST_DIR = "data/test"
BATCH_SIZE = 16
NUM_EPOCHS = 1
LEARNING_RATE = 1e-3
MODEL_SAVE_PATH = "checkpoints/best_model.pt"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ========== 创建文件夹 ==========
os.makedirs("checkpoints", exist_ok=True)
os.makedirs("data/processed", exist_ok=True)

# 1. 读取 + 预处理数据
train_raw, test_raw = concat_csv_files(TRAIN_DIR), concat_csv_files(TEST_DIR)
train_df, test_df = preprocess(train_raw, test_raw, save_dir="data")

train_dataset = TicketDataset(train_df)
test_dataset = TicketDataset(test_df)

# 2. DataLoader
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
val_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

# ========== 构建模型 ==========
num_numeric = train_dataset.numeric_feats.shape[1]
num_film_ids = len(train_dataset.encoder.classes_)

model = PricePredictor(
    num_numeric_features=num_numeric,
    num_film_ids=num_film_ids,
    film_emb_dim=16,
    hidden_dims=[128, 64],
    dropout=0.2
).to(DEVICE)

model_config = {
    "num_numeric_features": num_numeric,
    "num_film_ids": num_film_ids,
    "film_emb_dim": 16,
    "hidden_dims": [128, 64],
    "dropout": 0.2
}
with open("checkpoints/model_config.json", "w") as f:
    json.dump(model_config, f)

# ========== 损失函数与优化器 ==========
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

# ========== 训练循环 ==========
print("🚀 开始训练...")
best_val_loss = float("inf")

for epoch in range(1, NUM_EPOCHS + 1):
    model.train()
    total_train_loss = 0.0

    print(f"\n🎯 Epoch {epoch:02d}")
    train_bar = tqdm(train_loader, desc="🟢 Training", leave=False)
    for x_numeric, x_film_id, y in train_bar:
        x_numeric, x_film_id, y = x_numeric.to(DEVICE), x_film_id.to(DEVICE), y.to(DEVICE)

        optimizer.zero_grad()
        output = model(x_numeric, x_film_id)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()

        total_train_loss += loss.item() * x_numeric.size(0)
        train_bar.set_postfix(loss=loss.item())

    avg_train_loss = total_train_loss / len(train_loader.dataset)

    # 验证阶段
    model.eval()
    total_val_loss = 0.0
    val_bar = tqdm(val_loader, desc="🔵 Validating", leave=False)

    with torch.no_grad():
        for x_numeric, x_film_id, y in val_bar:
            x_numeric, x_film_id, y = x_numeric.to(DEVICE), x_film_id.to(DEVICE), y.to(DEVICE)
            output = model(x_numeric, x_film_id)
            loss = criterion(output, y)
            total_val_loss += loss.item() * x_numeric.size(0)
            val_bar.set_postfix(loss=loss.item())

    avg_val_loss = total_val_loss / len(val_loader.dataset)

    print(f"📊 Epoch {epoch:02d} | Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")

    # 保存最优模型
    if avg_val_loss < best_val_loss:
        best_val_loss = avg_val_loss
        torch.save(model.state_dict(), MODEL_SAVE_PATH)
        print(f"✅ 模型已保存：{MODEL_SAVE_PATH}（验证损失：{best_val_loss:.4f}）")